# -*- coding: utf-8 -*-
# @File: backend => websocket_persistent_test.py
# @CreateTime: 2025/1/27 
# @Author: EvanZhang
# @Info: WebSocket持久连接测试
import websockets
import asyncio
import json
import logging
from datetime import datetime

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class WebSocketClient:
    def __init__(self, uri, client_id):
        self.uri = uri
        self.client_id = client_id
        self.websocket = None
        self.running = False
        self.reconnect_interval = 5  # 重连间隔（秒）
        self.max_reconnect_attempts = 10
        self.reconnect_attempts = 0

    async def connect(self):
        """连接到WebSocket服务器"""
        try:
            self.websocket = await websockets.connect(self.uri)
            self.running = True
            self.reconnect_attempts = 0
            logger.info(f"Connected to {self.uri}")
            return True
        except Exception as e:
            logger.error(f"Failed to connect: {e}")
            return False

    async def disconnect(self):
        """断开WebSocket连接"""
        self.running = False
        if self.websocket:
            await self.websocket.close()
            logger.info("Disconnected from WebSocket")

    async def send_message(self, message_type="message", content=None):
        """发送消息到服务器"""
        if not self.websocket:
            logger.error("WebSocket not connected")
            return

        try:
            message = {
                "type": message_type,
                "content": content or f"Hello from client {self.client_id} at {datetime.now()}",
                "timestamp": asyncio.get_event_loop().time(),
                "client_id": self.client_id
            }
            await self.websocket.send(json.dumps(message))
            logger.info(f"Sent: {message}")
        except Exception as e:
            logger.error(f"Failed to send message: {e}")

    async def send_ping(self):
        """发送心跳包"""
        await self.send_message("ping", "ping")

    async def listen(self):
        """监听服务器消息"""
        try:
            while self.running:
                try:
                    # 设置接收超时
                    message = await asyncio.wait_for(self.websocket.recv(), timeout=70.0)
                    
                    try:
                        data = json.loads(message)
                        message_type = data.get("type", "unknown")
                        
                        if message_type == "heartbeat":
                            logger.debug(f"Received heartbeat: {data}")
                            # 响应心跳
                            await self.send_message("pong", "pong")
                        elif message_type == "pong":
                            logger.debug(f"Received pong: {data}")
                        elif message_type == "connection":
                            logger.info(f"Connection confirmed: {data['message']}")
                        elif message_type in ["user_joined", "user_left"]:
                            logger.info(f"User event: {data['message']}")
                        elif message_type in ["echo", "broadcast"]:
                            logger.info(f"Message received: {data['message']}")
                        else:
                            logger.info(f"Received: {data}")
                            
                    except json.JSONDecodeError:
                        logger.info(f"Received text: {message}")
                        
                except asyncio.TimeoutError:
                    logger.warning("No message received in 70 seconds, sending ping")
                    await self.send_ping()
                except websockets.exceptions.ConnectionClosed:
                    logger.warning("Connection closed by server")
                    break
                except Exception as e:
                    logger.error(f"Error receiving message: {e}")
                    break
                    
        except Exception as e:
            logger.error(f"Listen error: {e}")
        finally:
            self.running = False

    async def auto_reconnect(self):
        """自动重连机制"""
        while self.reconnect_attempts < self.max_reconnect_attempts:
            if not self.running:
                logger.info(f"Attempting to reconnect... (attempt {self.reconnect_attempts + 1})")
                if await self.connect():
                    logger.info("Reconnected successfully")
                    return True
                else:
                    self.reconnect_attempts += 1
                    await asyncio.sleep(self.reconnect_interval)
            else:
                await asyncio.sleep(1)
        
        logger.error("Max reconnection attempts reached")
        return False

    async def run_with_reconnect(self):
        """运行客户端并支持自动重连"""
        while True:
            if await self.connect():
                # 启动监听任务
                listen_task = asyncio.create_task(self.listen())
                
                # 启动定期发送消息任务（可选）
                async def send_periodic_messages():
                    await asyncio.sleep(5)  # 等待连接稳定
                    count = 0
                    while self.running:
                        await asyncio.sleep(30)  # 每30秒发送一条消息
                        if self.running:
                            count += 1
                            await self.send_message("message", f"Periodic message #{count}")
                
                message_task = asyncio.create_task(send_periodic_messages())
                
                # 等待任务完成
                try:
                    await asyncio.gather(listen_task, message_task)
                except Exception as e:
                    logger.error(f"Task error: {e}")
                finally:
                    listen_task.cancel()
                    message_task.cancel()
                    await self.disconnect()
            
            # 尝试重连
            if not await self.auto_reconnect():
                break

    async def run_interactive(self):
        """交互式运行模式"""
        if not await self.connect():
            return
        
        # 启动监听任务
        listen_task = asyncio.create_task(self.listen())
        
        try:
            print("WebSocket client started. Type 'quit' to exit.")
            print("Commands:")
            print("  - Type any message to send")
            print("  - 'ping' to send ping")
            print("  - 'quit' to exit")
            
            while self.running:
                try:
                    # 模拟用户输入（在实际应用中可以使用input()）
                    await asyncio.sleep(10)
                    await self.send_message("message", "Interactive test message")
                    
                except KeyboardInterrupt:
                    break
                    
        finally:
            listen_task.cancel()
            await self.disconnect()


async def test_persistent_connection():
    """测试持久连接"""
    client = WebSocketClient('ws://127.0.0.1:8000/api/ws/1', 1)
    
    try:
        # 运行带重连的客户端
        await client.run_with_reconnect()
    except KeyboardInterrupt:
        logger.info("Interrupted by user")
    finally:
        await client.disconnect()


async def test_multiple_clients():
    """测试多个客户端连接"""
    clients = []
    tasks = []
    
    try:
        # 创建多个客户端
        for i in range(1, 4):  # 创建3个客户端
            client = WebSocketClient(f'ws://127.0.0.1:8000/api/ws/{i}', i)
            clients.append(client)
            
            # 为每个客户端创建任务
            task = asyncio.create_task(client.run_with_reconnect())
            tasks.append(task)
            
            await asyncio.sleep(2)  # 间隔2秒创建下一个客户端
        
        # 等待所有任务
        await asyncio.gather(*tasks)
        
    except KeyboardInterrupt:
        logger.info("Interrupted by user")
    finally:
        # 清理所有客户端
        for client in clients:
            await client.disconnect()
        for task in tasks:
            task.cancel()


if __name__ == '__main__':
    print("WebSocket Persistent Connection Test")
    print("1. Test single persistent connection")
    print("2. Test multiple clients")
    
    choice = input("Choose test (1 or 2): ").strip()
    
    if choice == "1":
        asyncio.run(test_persistent_connection())
    elif choice == "2":
        asyncio.run(test_multiple_clients())
    else:
        print("Invalid choice, running single client test")
        asyncio.run(test_persistent_connection())